import matplotlib.pyplot as plt
import numpy as np
import os
from datetime import datetime


# 从文件中读取损失数据
def read_loss_data(file_path):
    losses = []
    with open(file_path, "r") as file:
        for line in file:
            if "loss in epoch" in line:
                # 提取损失值
                loss = float(line.strip().split(": ")[-1])
                losses.append(loss)
    return losses


# 创建输出目录
output_dir = "ml-1m_default"
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 使用当前时间创建文件名
current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
output_file = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    f"performance_comparison_{current_time}.png",
)


# 读取数据
# TODO 这里需要改地址
infini_transformer_losses = read_loss_data(
    "/Users/zzq/Developer/python/SAM/reproduction/infini-transformer/myCode/SASRec-in-InfiniTransformer/ml-1m_default/infiniTransformer_20250102_205318.txt"
)

routing_transformer_losses = read_loss_data(
    "/Users/zzq/Developer/python/SAM/sasrec_ex/myCode/ml-1m_output/routingTransformer_20250102_144322.txt"
)

transformer_losses = read_loss_data(
    "/Users/zzq/Developer/python/SAM/sasrec_ex/myCode/ml-1m_output/transformer-output.txt"
)

# 设置matplotlib样式参数
plt.rcParams["font.family"] = "Times New Roman"
plt.rcParams["axes.grid"] = True
plt.rcParams["grid.alpha"] = 0.3
plt.rcParams["grid.linestyle"] = "--"
plt.rcParams["axes.labelsize"] = 12
plt.rcParams["axes.titlesize"] = 16
plt.rcParams["figure.figsize"] = [16, 24]
plt.rcParams["legend.fontsize"] = 10
plt.rcParams["axes.spines.top"] = False
plt.rcParams["axes.spines.right"] = False
plt.rcParams["figure.subplot.top"] = 0.95
plt.rcParams["figure.subplot.bottom"] = 0.07
plt.rcParams["figure.subplot.hspace"] = 0.5

# 创建图表
plt.figure()

# 设置专业的配色方案
colors = ["#3498db", "#e74c3c", "#2ecc71", "#f1c40f"]

# 绘制损失对比图
plt.subplot(3, 1, 1)
plt.plot(
    infini_transformer_losses,
    label="Infini Transformer",
    linewidth=2.5,
    linestyle="-",
    color=colors[0],
)
plt.plot(
    routing_transformer_losses,
    label="Routing Transformer",
    linewidth=2.5,
    linestyle="--",
    color=colors[1],
)
plt.plot(
    transformer_losses,
    label="Simple Transformer",
    linewidth=2.5,
    linestyle="-.",
    color=colors[2],
)
plt.xlabel("Iteration", fontsize=12, fontweight="bold", labelpad=15)
plt.ylabel("Loss", fontsize=12, fontweight="bold", labelpad=10)
plt.title("Loss Comparison", fontsize=16, fontweight="bold", pad=20)

# 将图例移到图表外部右侧
plt.legend(
    loc="center left",
    fontsize=10,
    ncol=1,
    bbox_to_anchor=(1.02, 0.5),
)

# 运行时间对比图
plt.subplot(3, 1, 2)

# 创建三个空的散点图用于图例
plt.scatter([], [], color=colors[0], alpha=0.85, label="Infini Transformer")
plt.scatter([], [], color=colors[1], alpha=0.85, label="Routing Transformer")
plt.scatter([], [], color=colors[2], alpha=0.85, label="Simple Transformer")

# 绘制柱状图
valid_time_values = [1142.217348, 1407.427627, 1796.537418]
bars = plt.bar(
    ["Infini\nTransformer", "Routing\nTransformer", "Simple\nTransformer"],
    valid_time_values,
    color=colors[:3],
    alpha=0.85,
    width=0.5,  # 减小柱子宽度
)
plt.ylabel("Time (seconds)", fontsize=12, fontweight="bold", labelpad=10)
plt.title("Running Time Comparison", fontsize=16, fontweight="bold", pad=20)

# 调整y轴的范围，为标题腾出空间
plt.ylim(0, max(valid_time_values) * 1.2)

# 添加数值标签
for bar in bars:
    height = bar.get_height()
    plt.text(
        bar.get_x() + bar.get_width() / 2.0,
        height + max(valid_time_values) * 0.02,  # 调整标签位置
        f"{height:.1f}s",
        ha="center",
        va="bottom",
        fontsize=10,
    )

# 将图例移到图表外部右侧
plt.legend(
    loc="center left",
    fontsize=10,
    ncol=1,
    bbox_to_anchor=(1.02, 0.5),
)

# 准确性指标对比图
plt.subplot(3, 1, 3)
bar_width = 0.18
index = np.arange(3)

valid_ndcg_values = [0.6089, 0.5774, 0.6117]
valid_hr_values = [0.8406, 0.8131, 0.8455]
test_ndcg_values = [0.5835, 0.5536, 0.5838]
test_hr_values = [0.8157, 0.7871, 0.8157]

# 创建柱状图并存储返回的容器对象
rects1 = plt.bar(
    index - bar_width * 1.5,
    valid_ndcg_values,
    bar_width,
    label="Valid NDCG@10",
    color=colors[0],
    alpha=0.85,
)
rects2 = plt.bar(
    index - bar_width * 0.5,
    valid_hr_values,
    bar_width,
    label="Valid HR@10",
    color=colors[1],
    alpha=0.85,
)
rects3 = plt.bar(
    index + bar_width * 0.5,
    test_ndcg_values,
    bar_width,
    label="Test NDCG@10",
    color=colors[2],
    alpha=0.85,
)
rects4 = plt.bar(
    index + bar_width * 1.5,
    test_hr_values,
    bar_width,
    label="Test HR@10",
    color=colors[3],
    alpha=0.85,
)


# 为每个柱子添加数值标签
def autolabel(rects):
    for rect in rects:
        height = rect.get_height()
        plt.text(
            rect.get_x() + rect.get_width() / 2.0,
            height + 0.015,  # 调整标签位置
            f"{height:.3f}",
            ha="center",
            va="bottom",
            fontsize=8,
            rotation=0,
        )


autolabel(rects1)
autolabel(rects2)
autolabel(rects3)
autolabel(rects4)

plt.xlabel("Model", fontsize=12, fontweight="bold", labelpad=10)
plt.ylabel("Score", fontsize=12, fontweight="bold", labelpad=10)
plt.title("Accuracy Metrics Comparison", fontsize=16, fontweight="bold", pad=20)

# 调整y轴范围，为数值标签腾出空间
plt.ylim(0, max(max(valid_hr_values), max(test_hr_values)) * 1.15)

plt.xticks(
    index,
    ["Infini\nTransformer", "Routing\nTransformer", "Simple\nTransformer"],
    fontsize=10,
)

# 将图例移到图表外部右侧
plt.legend(
    loc="center left",
    fontsize=10,
    ncol=1,
    bbox_to_anchor=(1.02, 0.5),
)

# 统一调整所有子图的右边距
plt.subplots_adjust(right=0.85)

# 保存图表
plt.savefig(output_file, bbox_inches="tight", pad_inches=0.2)
plt.show()
